import os
import random
import sys
import time
import requests
import argparse
from queue import Queue
from threading import Thread
from lib.controller.checklist import init
from lib.core.addressExploit import data_write
from lib.core.statusCheck import alive_process
from lib.core.addressProcess import process_file, Address_ip
from lib.fileclass.fileProcess import FileHandler, RequestHandler
from colorama import init

# 初始化 colorama
init()

def get_random_color():
    # 返回随机的 ANSI 颜色代码
    colors = [
        '\033[1;31m',  # 红色
        '\033[1;32m',  # 绿色
        '\033[1;33m',  # 黄色
        '\033[1;34m',  # 蓝色
        '\033[1;35m',  # 紫色
        '\033[1;36m',  # 青色
        '\033[1;37m',  # 白色
    ]
    return random.choice(colors)

def print_pattern():
    pattern = """
     {color1}__     {color2}______{color1}  _    _   _      {color3}____{color1}   {color2}_____ {color1}______\033[0m
    {color1}/ \\   / {color2}/  _ \\{color1}| |  | | | |    / {color3}__ \\{color1} / {color2}_____{color1}|  ____|\033[0m
   {color1}/  /\\_/ /| {color2}|_) \\{color1}| |  | | | |   | {color3}|  | |{color1} | (___ | {color2}|__  \033[0m
  {color1}/  /    / | {color2}__/{color1}| |  | | | |   | {color3}|  | |{color1}\\___\\| {color2}|__| \033[0m   
 {color1}/  /    /  | {color2}|{color1}| |__|{color1}| |___| {color3}|__{color1}|{color2}|____\033[0m
/_/     /   |_|    \\____/  {color2}|_____/|______|\033[0m

         {color1}exploitsee by 00101\033[0m
    """
    # 随机选择颜色
    color1 = get_random_color()
    color2 = get_random_color()
    color3 = get_random_color()

    # 使用格式化字符串插入随机颜色
    random_pattern = pattern.format(color1=color1, color2=color2, color3=color3)

    # 打印随机颜色图案
    print(random_pattern)

def finalcheck(address_all, queue, consumer_thread):
    if not queue.empty():
        print(f"警告：队列中仍有 {queue.qsize()} 个任务未完成。")
        return False

    if consumer_thread.is_alive():
        print("警告：消费者线程未安全结束。")
        return False

    if len(address_all.get_addresses()) == 0:
        print("警告：地址处理未完全完成。")
        time.sleep(1)
        return False
    return True

def cmdline():
    parser = argparse.ArgumentParser(description="指定文件路径")
    parser.add_argument('-f', '--file', type=str, help='指定文件路径', required=True)
    args = parser.parse_args()
    return args.file

def start():
    file_path = cmdline()
    print_pattern()

    address_all = Address_ip()
    file_handler = FileHandler()
    request_handler = RequestHandler()

    queue = Queue()  #创建生产者消费者队列
    consumer_thread = Thread(target=data_write, args=(queue,))
    consumer_thread.start()
    session = requests.Session()

    process_file(file_path, address_all)
    print("地址全部处理成功...")

    for addr in address_all.get_addresses():
        alive_process(addr, queue, file_handler, request_handler, session)

    print("核心代码完成运行")

    file_handler.flush_cache()
    request_handler.flush_cache()
    queue.join()
    queue.put(None)
    consumer_thread.join()
    session.close()

    if not finalcheck(address_all, queue, consumer_thread):
        print("检查失败，程序存在问题。")
    else:
        print("所有检查通过，程序执行完毕。")


def main():
    init()
    start()


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        # 捕获用户Ctrl+C 导致的中断
        print("\n程序已被用户中断，正在退出...")
        os._exit(1)
    except Exception as e:
        # 捕获除 KeyboardInterrupt 以外的所有异常
        print(e)
    finally:
        sys.exit(1)
